#!/usr/bin/env python3
"""
Script: run_batch.py
Function: Batch process directories dedx_m2 to dedx_m24
          by copying scripts from dedx_m23 and running simulations
Author: Converted from bash script
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path


def main():
    # Colors for output (Windows compatible)
    RED = "\033[0;31m"
    GREEN = "\033[0;32m"
    YELLOW = "\033[1;33m"
    BLUE = "\033[0;34m"
    NC = "\033[0m"  # No Color

    # Source directory (contains files to copy)
    source_dir = "dedx_m23"

    # Check if source directory exists
    if not os.path.exists(source_dir):
        print(f"{RED}Error: Source directory {source_dir} does not exist{NC}")
        sys.exit(1)

    # Check if required source files exist
    required_files = ["run_rspw.py", "run_ionization.py"]
    for file in required_files:
        if not os.path.exists(os.path.join(source_dir, file)):
            print(f"{RED}Error: {os.path.join(source_dir, file)} not found in {source_dir}{NC}")
            sys.exit(1)

    print("Processing directories...")

    # Process directories dedx_m2 to dedx_m24
    for i in range(2, 25):
        dir_name = f"dedx_m{i}"
        print(f"Processing directory: {dir_name}")

        # Check if directory exists
        if not os.path.exists(dir_name):
            print(f"  {YELLOW}Warning: Directory {dir_name} does not exist, skipping{NC}")
            continue

        # Change to the directory
        original_dir = os.getcwd()
        try:
            os.chdir(dir_name)

            # Remove results and results1 subdirectories if they exist
            for results_dir in ["results", "results1"]:
                if os.path.exists(results_dir):
                    print(f"  Deleting subdirectory {results_dir}")
                    shutil.rmtree(results_dir)

            # Copy files from dedx_m23
            print(f"  Copying run_rspw.py and run_ionization.py from ../{source_dir}")
            for file in required_files:
                src_file = os.path.join("..", source_dir, file)
                dst_file = file
                shutil.copy2(src_file, dst_file)

            # Check if run.py exists and execute it
            if os.path.exists("run.py"):
                print("  Executing run.py")
                try:
                    subprocess.run([sys.executable, "run.py"], check=True)
                except subprocess.CalledProcessError as e:
                    print(f"  {RED}Error executing run.py: {e}{NC}")
            else:
                print("  Warning: run.py does not exist, skipping execution")

            # Return to original directory
            os.chdir(original_dir)
            print(f"  Completed processing {dir_name}")
            print()

        except Exception as e:
            print(f"  {RED}Error processing {dir_name}: {e}{NC}")
            os.chdir(original_dir)

    print("All directories processed!")


if __name__ == "__main__":
    main()
